// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::{hash_map, BTreeMap, BTreeSet, HashMap}, fmt::Display};

use serde::{Deserialize, Serialize};
use strum_macros::EnumCount;

use crate::{config::SpecsConfig, graph::DepKind, Arch};

use super::{overlay::OverlayConfig, SystemConfigError, TripleEnv};

#[derive(Debug, Serialize, Deserialize)]
pub struct ToolchainsConfig(pub BTreeMap<String, ToolchainConfig>);

#[derive(Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct ToolchainConfig {
    #[serde(flatten)]
    pub triple_envs: ToolchainTripleEnvsConfig,
    #[serde(flatten, default, skip_serializing_if = "HashMap::is_empty")]
    pub sub_toolchains: HashMap<SubToolchainKind, SubToolchainConfig>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub overlay: Option<OverlayConfig>,
    #[serde(default, skip_serializing_if = "SpecsConfig::is_empty")]
    pub temp_packages: SpecsConfig,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub extra_dep_edges: Vec<DepRelation>,
}

#[derive(Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct SubToolchainConfig {
    #[serde(flatten)]
    pub envs: ToolchainEnvsConfig,
    #[serde(default, skip_serializing_if = "SpecsConfig::is_empty")]
    pub packages: SpecsConfig,
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub cross_packages: HashMap<Arch, SpecsConfig>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, EnumCount)]
#[repr(u8)]
#[serde(rename_all = "kebab-case")]
pub enum SubToolchainKind {
    Basic,
    Cpp,
    Fortran,
    Rust,
    Meson,
}

#[derive(Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct ToolchainEnvsConfig {
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub envs: BTreeMap<String, String>,
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub cross_envs: HashMap<Arch, BTreeMap<String, String>>,
}

#[derive(Debug, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct ToolchainTripleEnvsConfig {
    #[serde(default, skip_serializing_if = "TripleEnv::is_default")]
    pub triple_env: TripleEnv,
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub cross_triple_envs: HashMap<Arch, TripleEnv>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DepRelation {
    pub source: String,
    pub target: String,
    pub kind: DepKind,
}

impl Display for SubToolchainKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

impl TryFrom<u8> for SubToolchainKind {
    type Error = u8;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        if value < <Self as strum::EnumCount>::COUNT as u8 {
            Ok(unsafe { std::mem::transmute::<u8, Self>(value) })
        } else {
            Err(value)
        }
    }
}

impl ToolchainsConfig {
    pub fn verify(&self) -> Result<(), SystemConfigError> {
        for (toolchain_name, toolchain_config) in self.0.iter() {
            toolchain_config.verify(toolchain_name)?;
        }
        Ok(())
    }
}

impl ToolchainConfig {
    fn verify(&self, toolchain_name: &String) -> Result<(), SystemConfigError> {
        let mut names = BTreeSet::new();
        for name in self.sub_toolchains.values().flat_map(|t| t.packages.0.keys())
            .chain(self.sub_toolchains.values().flat_map(|t| t.cross_packages.values().flat_map(|map| map.0.keys())))
            .chain(self.temp_packages.0.keys())
        {
            if !names.insert(name.clone()) {
                return Err(SystemConfigError::DuplicatedPkgName(toolchain_name.clone(), name.clone()));
            }
        }

        for extra_dep_edge in self.extra_dep_edges.iter() {
            if !self.contains_package(&extra_dep_edge.source) {
                return Err(SystemConfigError::EdgePkgnameNotFound(toolchain_name.clone(), extra_dep_edge.source.clone()));
            }
            if !self.contains_package(&extra_dep_edge.target) {
                return Err(SystemConfigError::EdgePkgnameNotFound(toolchain_name.clone(), extra_dep_edge.target.clone()));
            }
        }
        Ok(())
    }

    fn contains_package(&self, name: &String) -> bool {
        self.sub_toolchains.values().any(|t| t.contains_package(name))
        || self.temp_packages.0.contains_key(name)
    }

    pub fn sort_features(&mut self) {
        for spec in self.temp_packages.0.values_mut() {
            spec.features.sort();
        }
        for sub_toolchain in self.sub_toolchains.values_mut() {
            sub_toolchain.sort_features();
        }
    }
}

impl SubToolchainConfig {
    fn contains_package(&self, name: &String) -> bool {
        self.packages.0.contains_key(name)
        || self.cross_packages.values().any(|map| map.0.contains_key(name))
    }

    fn sort_features(&mut self) {
        for spec in self.packages.0.values_mut() {
            spec.features.sort();
        }
        for specs in self.cross_packages.values_mut() {
            for spec in specs.0.values_mut() {
                spec.features.sort();
            }
        }
    }
}

impl SubToolchainKind {
    pub fn as_str(&self) -> &'static str {
        match self {
            SubToolchainKind::Basic => "basic",
            SubToolchainKind::Cpp => "cpp",
            SubToolchainKind::Fortran => "fortran",
            SubToolchainKind::Rust => "rust",
            SubToolchainKind::Meson => "meson",
        }
    }

    pub fn is_default(&self) -> bool {
        matches!(self, SubToolchainKind::Basic | SubToolchainKind::Cpp)
    }
}

impl ToolchainEnvsConfig {
    pub fn get_envs(&self, arch: Arch) -> Option<&BTreeMap<String, String>> {
        if arch.is_host() {
            Some(&self.envs)
        } else {
            self.cross_envs.get(&arch)
        }
    }

    pub fn insert_envs<I>(&mut self, arch: Arch, envs: I)
    where I: Iterator<Item = (String, String)>
    {
        if arch.is_host() {
            self.envs.extend(envs);
        } else {
            match self.cross_envs.entry(arch) {
                hash_map::Entry::Occupied(mut occupied_entry) => {
                    occupied_entry.get_mut().extend(envs);
                },
                hash_map::Entry::Vacant(vacant_entry) => {
                    vacant_entry.insert(envs.collect());
                },
            }
        }
    }
}

impl ToolchainTripleEnvsConfig {
    pub fn get_env(&self, arch: Arch) -> TripleEnv {
        if arch.is_host() {
            self.triple_env
        } else {
            self.cross_triple_envs.get(&arch).copied().unwrap_or_default()
        }
    }

    pub fn make_triple(&self, arch: Arch) -> String {
        self.get_env(arch).to_triple(arch)
    }

    pub fn make_normalized_triple(&self, arch: Arch) -> String {
        self.get_env(arch).to_normalized_triple(arch)
    }
}
